#include "VulkanApplication.h"
#include "BufferManager.h"
#include "../glm/glm.hpp"
#include "../glm/gtc/matrix_transform.hpp"
#include <vector>
#include <iostream>
#include <fstream>
#include <cstring>
#include <cmath>

VkSampler createTextureSampler(VulkanApplication& app)
{
    VkPhysicalDeviceProperties properties;
    vkGetPhysicalDeviceProperties(app.getPhysicalDevice(), &properties);

    VkSamplerCreateInfo samplerInfo
    {
        VK_STRUCTURE_TYPE_SAMPLER_CREATE_INFO,
        nullptr, 0,
        VK_FILTER_LINEAR, VK_FILTER_LINEAR,
        VK_SAMPLER_MIPMAP_MODE_LINEAR,
        VK_SAMPLER_ADDRESS_MODE_REPEAT,
        VK_SAMPLER_ADDRESS_MODE_REPEAT,
        VK_SAMPLER_ADDRESS_MODE_REPEAT,
        0.0f,  // lodBias
        VK_TRUE, properties.limits.maxSamplerAnisotropy,  // anisotropy
        VK_FALSE, VK_COMPARE_OP_ALWAYS,  // compare op
        0.0f, 0.0f,  // minLod, maxLod
        VK_BORDER_COLOR_INT_OPAQUE_BLACK,
        VK_FALSE  // unnormalized coordinates
    };
    
    VkSampler texSampler = nullptr;
    VkResult result = vkCreateSampler(app.getDevice(), &samplerInfo, nullptr, &texSampler);
    if (result != VK_SUCCESS)
    {
        printf("Failed to create texture sampler.\n");
        return nullptr;
    }
    return texSampler;
}

VkImageView createImageView(VulkanApplication& app, VkImage image)
{
    VkImageViewCreateInfo viewInfo
    {
        VK_STRUCTURE_TYPE_IMAGE_VIEW_CREATE_INFO,
        nullptr, 0,
        image, VK_IMAGE_VIEW_TYPE_2D,
        VK_FORMAT_R8G8B8A8_SRGB,
        { VK_COMPONENT_SWIZZLE_IDENTITY, VK_COMPONENT_SWIZZLE_IDENTITY,
          VK_COMPONENT_SWIZZLE_IDENTITY, VK_COMPONENT_SWIZZLE_IDENTITY },
        { VK_IMAGE_ASPECT_COLOR_BIT, 0, 1, 0, 1 }
    };

    VkImageView imageView = nullptr;
    VkResult result = vkCreateImageView(app.getDevice(), &viewInfo, nullptr, &imageView);
    if (result != VK_SUCCESS)
    {
        printf("Failed to create texture image view.\n");
        return nullptr;
    }
    return imageView;
}

int main(int argc, char** argv)
{
    VulkanApplication app("c14_animation", 1024, 768);
    if (!app.initialized())
    {
        printf("Failed to initialize VulkanApplication.\n");
        return 1;
    }

    // Create pipeline, vertices and indices
    BufferManager manager(&app);
    if (!manager.createDescriptorSet() ||
        !manager.prepare("draw_texture.vert", "draw_texture.frag"))
    {
        printf("Failed to create pipeline.\n");
        return 1;
    }

    // Create vertex & index buffer
    std::vector<Vertex> vertices;
    vertices.push_back({ -0.5f, -0.5f, 0.0f, 0.0f, 1.0f, 0.0f, /*uv*/ 0.0f, 0.0f });
    vertices.push_back({ 0.5f, -0.5f, 0.0f, 0.0f, 0.0f, 1.0f, /*uv*/ 1.0f, 0.0f });
    vertices.push_back({ 0.5f, 0.5f, 0.0f, 1.0f, 0.0f, 0.0f, /*uv*/ 1.0f, 1.0f });
    vertices.push_back({ -0.5f, 0.5f, 0.0f, 1.0f, 1.0f, 0.0f, /*uv*/ 0.0f, 1.0f });

    std::vector<unsigned int> indices;
    indices.push_back(0); indices.push_back(1); indices.push_back(2);
    indices.push_back(0); indices.push_back(2); indices.push_back(3);
    
    struct UniformBufferObject
    {
        glm::mat4x4 modelview;
        glm::mat4x4 projection;
    };
    UniformBufferObject ubo;

    VkBuffer vertexBuffer = manager.createVertexBuffer("vertices", vertices);
    VkBuffer indexBuffer = manager.createIndexBuffer("indices", indices);
    VkBuffer uboBuffer = manager.createUniformBuffer("ubo", sizeof(UniformBufferObject));

    // Create mapping of vertices and ubo
    void* vertexData = manager.mapBufferMemory("vertices", vertices.size() * sizeof(Vertex));
    void* uboData = manager.mapBufferMemory("ubo", sizeof(UniformBufferObject));
    
    // Create texture image
    VkImage textureImage = manager.createTextureImage("image", "lingjing.jpg");
    VkImageView textureImageView = createImageView(app, textureImage);
    VkSampler textureSampler = createTextureSampler(app);

    VkDescriptorImageInfo imageInfo
    {
        textureSampler, textureImageView,
        VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL
    };

    VkWriteDescriptorSet descriptorWrite
    {
        VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET,
        nullptr,
        manager.getDescriptorSet(), 1,  // dstSet, dstBinding
        0, 1,              // dstArrayElement, descriptorCount
        VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER,
        &imageInfo, nullptr, nullptr
    };
    vkUpdateDescriptorSets(app.getDevice(), 1, &descriptorWrite, 0, nullptr);

    // Main loop
    int width = 1024, height = 768;
    VkFence waitFence = app.getOrCreateFence("CommandFence");
    VkSemaphore waitNextImage = app.getOrCreateSemaphore("WaitNextImage");
    VkSemaphore waitSubmission = app.getOrCreateSemaphore("WaitSubmission");
    VkCommandBuffer commandBuffer = app.getCommandBuffer();

    float angleZ = 0.0f, aspectRatio = 0.0f;
    while (!app.shouldClose())
    {
        vkWaitForFences(app.getDevice(), 1, &waitFence, VK_TRUE, UINT64_MAX);
        vkResetFences(app.getDevice(), 1, &waitFence);

        uint32_t imageIndex = 0;
        VkResult result = vkAcquireNextImageKHR(app.getDevice(), app.getSwapchain(), UINT64_MAX,
                                                waitNextImage, VK_NULL_HANDLE, &imageIndex);
        switch (result)
        {
        case VK_SUBOPTIMAL_KHR:
        case VK_ERROR_OUT_OF_DATE_KHR:
            imageIndex = app.recreateSwapchain(width, height, waitNextImage); break;
        case VK_SUCCESS:
            break;
        default:
            printf("Next image: %d (RESULT = %d)\n", imageIndex, result);
            app.setClosed(); break;
        }

        for (size_t i = 0; i < vertices.size(); ++i)
        {
            Vertex& vec = vertices[i];
            vec.u += 0.0001f; vec.v += 0.0001f;
        }
        memcpy(vertexData, vertices.data(), vertices.size() * sizeof(Vertex));

        angleZ += 0.01f; aspectRatio = (float)width / (float)height;
        glm::mat4 rotation = glm::rotate(glm::mat4(1.0f), glm::radians(angleZ), glm::vec3(0.0f, 0.0f, 1.0f));
        glm::vec3 up = glm::vec3(rotation * glm::vec4(0.0f, 1.0f, 0.0f, 1.0f));
        ubo.modelview = glm::lookAt(glm::vec3(0.0f, 0.0f, 2.0f), glm::vec3(), up);
        ubo.projection = glm::perspective(45.0f, aspectRatio, 0.1f, 100.0f);
        memcpy(uboData, &ubo, sizeof(UniformBufferObject));

        app.beginRendering(imageIndex, width, height);
        {
            VkBuffer vertexBuffers[] = { vertexBuffer };
            VkDeviceSize offsets[] = { 0 };
            VkDescriptorSet descrptorSet = manager.getDescriptorSet();
            vkCmdBindDescriptorSets(commandBuffer, VK_PIPELINE_BIND_POINT_GRAPHICS,
                                    app.getPipelineLayout(), 0, 1, &descrptorSet, 0, nullptr);

            vkCmdBindVertexBuffers(commandBuffer, 0, 1, vertexBuffers, offsets);
            vkCmdBindIndexBuffer(commandBuffer, indexBuffer, 0, VK_INDEX_TYPE_UINT32);
            vkCmdDrawIndexed(commandBuffer, (uint32_t)indices.size(), 1, 0, 0, 0);
        }
        app.endRendering();

        app.submitAndPresent(waitNextImage, waitSubmission, waitFence, imageIndex);
        glfwPollEvents();
    }

    vkDeviceWaitIdle(app.getDevice());
    vkDestroySampler(app.getDevice(), textureSampler, nullptr);
    vkDestroyImageView(app.getDevice(), textureImageView, nullptr);
    manager.unmapBufferMemory("vertices");
    manager.unmapBufferMemory("ubo");
    manager.cleanup();
    return 0;
}